from .dalchemy_llm import DalchemyLLM
from typing import Dict
from .api_llm import AzureLLM,OpenAILLM,ZhipuLLM
import os
import logging
import importlib

LLM_REGISTRY: Dict[str, DalchemyLLM] = {}


def register_llm(name):
    def register_llm_fn(cls):
        assert issubclass(cls, DalchemyLLM)
        if cls in LLM_REGISTRY.values():
            raise ValueError("Cannot register duplicate llm ({})".format(name))
        LLM_REGISTRY[name] = cls
        return cls

    return register_llm_fn


def import_models(infer_dir, namespace):
    for file in os.listdir(infer_dir):
        path = os.path.join(infer_dir, file)
        if (
                not file.startswith("_")
                and not file.startswith(".")
                and (file.endswith(".py") or os.path.isdir(path))
        ):
            infer_name = file[: file.find(".py")] if file.endswith(".py") else file
            importlib.import_module(namespace + "." + infer_name)


# automatically import any Python files in the infer/ directory
infer_dir = os.path.dirname(__file__)
import_models(infer_dir, "dalchemy.llms")


def build_llm(name):
    return LLM_REGISTRY.get(name, DalchemyLLM)()

